package logger

import (
	"gin-skeleton/core"
	"github.com/sirupsen/logrus"
	"sync"
)

var (
	traceHook *TraceIdHook
)

type TraceIdHook struct {
	TraceId *string
	mu      sync.RWMutex
}

// SetTraceId 设置日志跟踪ID
func SetTraceId(traceId string) {

	if usageRotate {
		// 文件日志
		LogEntry = Get().WithField(core.RequestID, traceId)
	}

	// 标准控制台日志
	if traceHook == nil {
		traceHook = &TraceIdHook{}
		logrus.AddHook(traceHook)
	}
	traceHook.mu.Lock()
	defer traceHook.mu.Unlock()
	traceHook.TraceId = &traceId
}

func (h *TraceIdHook) Fire(entry *logrus.Entry) error {
	h.mu.RLock()
	defer h.mu.RUnlock()
	entry.Data[core.RequestID] = *h.TraceId
	return nil
}

func (h *TraceIdHook) Levels() []logrus.Level {
	return logrus.AllLevels
}
